import sys
import os
import cv2
import numpy as np
import torch
import time
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QPushButton, QLabel, QFileDialog, QComboBox, QProgressBar,
                             QScrollArea, QSplitter, QFrame, QMessageBox, QSlider)
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QSize, QTimer
from PyQt5.QtGui import QPixmap, QImage, QFont, QIcon


class YOLOModel:
    """
    YOLO模型加载和推理类 - 只加载本地模型
    """

    def __init__(self, model_path):
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        print(f"使用设备: {self.device}")

        try:
            # 仅加载本地模型
            if model_path and os.path.exists(model_path):
                print(f"正在加载本地模型: {model_path}")
                from ultralytics import YOLO
                self.model = YOLO(model_path)
                self.model.to(self.device)
                print("YOLO模型加载成功")
            else:
                print(f"错误: 指定的模型文件不存在: {model_path}")
                self.model = None
        except Exception as e:
            print(f"模型加载失败: {str(e)}")
            self.model = None

    def detect(self, image, conf_threshold=0.25):
        """
        对输入图像进行缺陷检测
        """
        if self.model is None:
            return None, "模型未成功加载"

        try:
            # 使用YOLO模型进行检测
            results = self.model.predict(image, conf=conf_threshold, verbose=False)[0]

            # 处理结果
            detections = []
            boxes = results.boxes

            if len(boxes) > 0:
                for i, box in enumerate(boxes):
                    # 获取边界框
                    x1, y1, x2, y2 = map(int, box.xyxy.cpu().numpy()[0])

                    # 获取置信度
                    confidence = float(box.conf.cpu().numpy()[0])

                    # 获取类别
                    class_id = int(box.cls.cpu().numpy()[0])
                    class_name = results.names[class_id]

                    detections.append({
                        'bbox': (x1, y1, x2, y2),
                        'class_id': class_id,
                        'class_name': class_name,
                        'confidence': confidence
                    })

            # 在图像上绘制检测结果
            result_image = image.copy()
            for det in detections:
                x1, y1, x2, y2 = det['bbox']
                class_name = det['class_name']
                confidence = det['confidence']

                # 不同类型的缺陷使用不同颜色
                color = (0, 0, 255)  # 默认红色
                if 'missing' in class_name.lower():
                    color = (0, 0, 255)  # 红色表示缺失
                elif 'short' in class_name.lower():
                    color = (255, 0, 0)  # 蓝色表示短路
                elif 'open' in class_name.lower():
                    color = (0, 255, 0)  # 绿色表示开路

                cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 2)
                label = f"{class_name}: {confidence:.2f}"
                cv2.putText(result_image, label, (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

            return result_image, detections

        except Exception as e:
            print(f"检测过程中出错: {str(e)}")
            return None, str(e)


class DetectionThread(QThread):
    """
    执行图像检测的线程类，避免GUI冻结
    """
    update_signal = pyqtSignal(object, object)
    error_signal = pyqtSignal(str)
    progress_signal = pyqtSignal(int)

    def __init__(self, model, image_path, conf_threshold, is_file=True):
        super().__init__()
        self.model = model
        self.image_path = image_path
        self.conf_threshold = conf_threshold
        self.is_file = is_file  # 是文件还是摄像头索引

    def run(self):
        try:
            self.progress_signal.emit(10)

            if self.is_file:
                # 读取图像文件
                image = cv2.imread(self.image_path)
                if image is None:
                    self.error_signal.emit("无法读取图像文件")
                    return
            else:
                # 此处不处理摄像头，只处理单张图像
                self.error_signal.emit("无效的图像源")
                return

            self.progress_signal.emit(30)
            # 执行检测
            result_image, detections = self.model.detect(image, self.conf_threshold)
            self.progress_signal.emit(90)

            if result_image is None:
                self.error_signal.emit(f"检测失败: {detections}")
                return

            self.progress_signal.emit(100)
            # 发送检测结果信号
            self.update_signal.emit(result_image, detections)

        except Exception as e:
            self.error_signal.emit(f"检测过程中出错: {str(e)}")


class VideoDetectionThread(QThread):
    """
    执行视频流检测的线程类
    """
    frame_signal = pyqtSignal(object)
    detection_signal = pyqtSignal(object)
    error_signal = pyqtSignal(str)
    fps_signal = pyqtSignal(float)

    def __init__(self, model, video_source, conf_threshold):
        super().__init__()
        self.model = model
        self.video_source = video_source  # 可以是摄像头索引或视频文件路径
        self.conf_threshold = conf_threshold
        self.running = False
        self.fps = 0

    def run(self):
        try:
            # 打开视频源
            cap = cv2.VideoCapture(self.video_source)
            if not cap.isOpened():
                self.error_signal.emit(f"无法打开视频源: {self.video_source}")
                return

            self.running = True
            frame_count = 0
            start_time = time.time()
            processing_time = 0

            while self.running:
                # 读取一帧
                ret, frame = cap.read()
                if not ret:
                    # 如果是视频文件，可以循环播放
                    if isinstance(self.video_source, str):
                        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
                        continue
                    else:
                        # 摄像头没有读到帧，可能断开了
                        self.error_signal.emit("视频流中断")
                        break

                # 计算FPS
                frame_count += 1
                if frame_count % 10 == 0:  # 每10帧更新一次FPS
                    end_time = time.time()
                    elapsed_time = end_time - start_time
                    self.fps = frame_count / elapsed_time if elapsed_time > 0 else 0
                    self.fps_signal.emit(self.fps)

                # 仅在处理完上一帧后才处理新帧，避免积压
                if processing_time <= 0 or time.time() - processing_time > 0.1:  # 限制处理频率
                    processing_time = time.time()

                    # 发送原始帧
                    self.frame_signal.emit(frame.copy())

                    # 执行检测
                    try:
                        result_image, detections = self.model.detect(frame, self.conf_threshold)
                        if result_image is not None:
                            # 在结果图像上添加FPS信息
                            cv2.putText(result_image, f"FPS: {self.fps:.1f}", (10, 30),
                                        cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)

                            # 发送检测结果
                            self.detection_signal.emit((result_image, detections))
                    except Exception as e:
                        print(f"帧处理错误: {str(e)}")

            # 释放资源
            cap.release()

        except Exception as e:
            self.error_signal.emit(f"视频处理错误: {str(e)}")

        finally:
            self.running = False

    def stop(self):
        """停止视频检测线程"""
        self.running = False
        self.wait(1000)  # 等待线程结束


class PCBDefectDetectionApp(QMainWindow):
    """
    电路板缺陷检测应用主窗口
    """

    def __init__(self):
        super().__init__()

        # 应用设置
        self.setWindowTitle("电路板缺陷检测系统 - 图像与视频流检测")
        self.setMinimumSize(1200, 800)

        # 初始化UI
        self.init_ui()

        # 初始化模型
        self.model = None
        self.init_model()

        # 状态变量
        self.current_image_path = None
        self.current_image = None
        self.detection_results = None
        self.video_thread = None  # 视频检测线程
        self.is_video_mode = False  # 是否处于视频模式

    def init_model(self):
        """初始化YOLO模型 - 仅使用本地模型"""
        try:
            self.status_label.setText("正在加载本地模型...")
            QApplication.processEvents()

            # 定义本地模型路径
            # local_model_path = 'E:/bishepcb/yolo11-circuit/yolo11-circuit/ultralytics-main/runs/detect/train2/weights/best.pt'
            # local_model_path = 'C:/Users/19002/Desktop/yolo11-circuit/yolo11n.pt'
            local_model_path = 'C:/Users/19002/Desktop/yolo11-circuit-research/runs/detect/train/weights/best.pt'
                                # C:\Users\19002\Desktop\yolo11-circuit-research\ultralytics\runs\detect\train\weights\best.pt
            # 检查模型文件是否存在
            if not os.path.exists(local_model_path):
                error_msg = f"找不到本地模型文件: {local_model_path}"
                self.status_label.setText(error_msg)
                QMessageBox.critical(self, "错误", f"{error_msg}\n请确保模型文件放在正确位置。")
                return

            # 加载本地模型
            self.model = YOLOModel(model_path=local_model_path)

            if self.model.model is None:
                error_msg = f"模型加载失败，请检查模型文件格式"
                self.status_label.setText(error_msg)
                QMessageBox.critical(self, "错误", error_msg)
                return

            self.status_label.setText("本地模型加载完成，可以开始检测")
        except Exception as e:
            error_msg = f"模型加载失败: {str(e)}"
            self.status_label.setText(error_msg)
            QMessageBox.critical(self, "错误", error_msg)

    def init_ui(self):
        """初始化用户界面"""
        # 创建中央部件
        central_widget = QWidget(self)
        self.setCentralWidget(central_widget)

        # 主布局
        main_layout = QVBoxLayout(central_widget)

        # 标题
        title_label = QLabel("电路板缺陷检测系统 (YOLOv11)")
        title_label.setAlignment(Qt.AlignCenter)
        title_label.setFont(QFont("Arial", 18, QFont.Bold))
        main_layout.addWidget(title_label)

        # 分割器
        splitter = QSplitter(Qt.Horizontal)
        main_layout.addWidget(splitter, 1)

        # === 左侧面板 ===
        left_panel = QWidget()
        left_layout = QVBoxLayout(left_panel)

        # 图像显示区域
        self.image_label = QLabel("请上传图像")
        self.image_label.setAlignment(Qt.AlignCenter)
        self.image_label.setMinimumSize(400, 400)
        self.image_label.setFrameShape(QFrame.Box)
        self.image_label.setStyleSheet("background-color: #f0f0f0;")

        # 创建滚动区域用于图像缩放和滚动
        scroll_area = QScrollArea()
        scroll_area.setWidget(self.image_label)
        scroll_area.setWidgetResizable(True)
        left_layout.addWidget(scroll_area)

        # 图像操作按钮
        buttons_layout = QHBoxLayout()

        self.load_button = QPushButton("上传图像")
        self.load_button.clicked.connect(self.load_image)
        buttons_layout.addWidget(self.load_button)

        self.detect_button = QPushButton("检测缺陷")
        self.detect_button.clicked.connect(self.detect_defects)
        self.detect_button.setEnabled(False)
        buttons_layout.addWidget(self.detect_button)

        self.save_button = QPushButton("保存结果")
        self.save_button.clicked.connect(self.save_results)
        self.save_button.setEnabled(False)
        buttons_layout.addWidget(self.save_button)

        left_layout.addLayout(buttons_layout)

        # 视频检测按钮
        video_buttons_layout = QHBoxLayout()

        self.camera_button = QPushButton("开启摄像头")
        self.camera_button.clicked.connect(lambda: self.start_video_detection(0))  # 0是默认摄像头索引
        video_buttons_layout.addWidget(self.camera_button)

        self.video_button = QPushButton("加载视频")
        self.video_button.clicked.connect(self.load_video)
        video_buttons_layout.addWidget(self.video_button)

        self.stop_video_button = QPushButton("停止检测")
        self.stop_video_button.clicked.connect(self.stop_video_detection)
        self.stop_video_button.setEnabled(False)
        video_buttons_layout.addWidget(self.stop_video_button)

        left_layout.addLayout(video_buttons_layout)

        # FPS显示
        self.fps_label = QLabel("FPS: 0.0")
        self.fps_label.setAlignment(Qt.AlignRight)
        left_layout.addWidget(self.fps_label)

        # 检测配置
        config_layout = QHBoxLayout()

        config_layout.addWidget(QLabel("置信度阈值:"))
        self.confidence_slider = QSlider(Qt.Horizontal)
        self.confidence_slider.setRange(1, 100)
        self.confidence_slider.setValue(25)  # 默认值0.25
        self.confidence_slider.setTickPosition(QSlider.TicksBelow)
        self.confidence_slider.setTickInterval(5)
        config_layout.addWidget(self.confidence_slider)

        self.threshold_label = QLabel("0.25")
        config_layout.addWidget(self.threshold_label)

        self.confidence_slider.valueChanged.connect(self.update_threshold_label)

        left_layout.addLayout(config_layout)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        self.progress_bar.setVisible(False)
        left_layout.addWidget(self.progress_bar)

        # === 右侧面板 ===
        right_panel = QWidget()
        right_layout = QVBoxLayout(right_panel)

        right_layout.addWidget(QLabel("检测结果:"))

        # 结果列表
        self.results_label = QLabel("未执行检测")
        self.results_label.setAlignment(Qt.AlignTop | Qt.AlignLeft)
        self.results_label.setWordWrap(True)
        self.results_label.setMinimumWidth(300)
        self.results_label.setFrameShape(QFrame.Box)
        self.results_label.setStyleSheet("background-color: white; padding: 5px;")

        results_scroll = QScrollArea()
        results_scroll.setWidget(self.results_label)
        results_scroll.setWidgetResizable(True)
        right_layout.addWidget(results_scroll, 1)

        # 添加面板到分割器
        splitter.addWidget(left_panel)
        splitter.addWidget(right_panel)
        splitter.setSizes([600, 400])  # 设置初始分割比例

        # 状态栏
        self.statusBar()
        self.status_label = QLabel("准备就绪")
        self.statusBar().addWidget(self.status_label)

    def update_threshold_label(self):
        """更新置信度阈值的显示"""
        value = self.confidence_slider.value() / 100
        self.threshold_label.setText(f"{value:.2f}")

    def load_image(self):
        """加载图像文件"""
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择图像", "",
            "图像文件 (*.jpg *.jpeg *.png *.bmp);;所有文件 (*)",
            options=options
        )

        if file_path:
            try:
                # 读取图像并显示
                self.current_image_path = file_path
                self.current_image = cv2.imread(file_path)
                self.display_image(self.current_image)

                # 更新UI状态
                self.detect_button.setEnabled(True)
                self.save_button.setEnabled(False)
                self.results_label.setText("已加载图像，等待检测")
                self.status_label.setText(f"已加载图像: {os.path.basename(file_path)}")
            except Exception as e:
                QMessageBox.critical(self, "错误", f"无法加载图像: {str(e)}")

    def display_image(self, cv_image):
        """在GUI中显示OpenCV图像"""
        if cv_image is None:
            return

        # 转换OpenCV的BGR图像到RGB
        rgb_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
        h, w, ch = rgb_image.shape

        # 转换为QImage
        qt_image = QImage(rgb_image.data, w, h, w * ch, QImage.Format_RGB888)

        # 获取显示区域大小
        display_size = self.image_label.size()

        # 如果图像太大，进行等比例缩放以适应显示区域
        pixmap = QPixmap.fromImage(qt_image)
        if pixmap.width() > display_size.width() or pixmap.height() > display_size.height():
            pixmap = pixmap.scaled(display_size, Qt.KeepAspectRatio, Qt.SmoothTransformation)

        # 设置图像
        self.image_label.setPixmap(pixmap)

        # 在视频模式下不改变Label大小，否则会导致闪烁
        if not self.is_video_mode:
            self.image_label.setMinimumSize(pixmap.size())

    def detect_defects(self):
        """检测电路板缺陷"""
        if self.current_image_path is None or self.current_image is None:
            QMessageBox.warning(self, "警告", "请先加载图像")
            return

        if self.model is None:
            QMessageBox.critical(self, "错误", "模型未成功加载")
            return

        # 从滑块获取置信度阈值
        conf_threshold = self.confidence_slider.value() / 100

        # 设置UI状态
        self.progress_bar.setValue(0)
        self.progress_bar.setVisible(True)
        self.detect_button.setEnabled(False)
        self.status_label.setText("正在检测...")

        try:
            # 创建并启动检测线程
            self.detection_thread = DetectionThread(self.model, self.current_image_path, conf_threshold)
            self.detection_thread.update_signal.connect(self.handle_detection_result)
            self.detection_thread.error_signal.connect(self.handle_detection_error)
            self.detection_thread.progress_signal.connect(self.progress_bar.setValue)
            self.detection_thread.start()
        except Exception as e:
            self.handle_detection_error(f"启动检测线程失败: {str(e)}")

    def handle_detection_result(self, result_image, detections):
        """处理检测结果"""
        # 更新UI
        self.display_image(result_image)
        self.detection_results = detections
        self.save_button.setEnabled(True)
        self.detect_button.setEnabled(True)
        self.progress_bar.setVisible(False)
        self.status_label.setText(f"检测完成，找到 {len(detections)} 个缺陷")

        # 显示检测结果
        if not detections:
            self.results_label.setText("未检测到缺陷")
        else:
            result_text = "<h3>检测到的缺陷:</h3>"
            for i, det in enumerate(detections):
                class_name = det['class_name']
                confidence = det['confidence']
                x1, y1, x2, y2 = det['bbox']
                result_text += f"<p><b>缺陷 {i + 1}:</b><br>"
                result_text += f"类型: {class_name}<br>"
                result_text += f"置信度: {confidence:.2f}<br>"
                result_text += f"位置: ({x1}, {y1}) 到 ({x2}, {y2})</p>"

            self.results_label.setText(result_text)

    def handle_detection_error(self, error_message):
        """处理检测过程中的错误"""
        self.progress_bar.setVisible(False)
        self.detect_button.setEnabled(True)
        self.status_label.setText("检测失败")
        QMessageBox.critical(self, "检测错误", error_message)

    def load_video(self):
        """加载视频文件"""
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getOpenFileName(
            self, "选择视频文件", "",
            "视频文件 (*.mp4 *.avi *.mov *.wmv);;所有文件 (*)",
            options=options
        )

        if file_path:
            self.start_video_detection(file_path)

    def start_video_detection(self, video_source):
        """启动视频检测

        Args:
            video_source: 可以是摄像头索引（整数）或视频文件路径（字符串）
        """
        if self.model is None or self.model.model is None:
            QMessageBox.critical(self, "错误", "模型未成功加载")
            return

        # 如果已经有视频线程在运行，先停止它
        if self.video_thread is not None and self.video_thread.running:
            self.stop_video_detection()

        # 显示视频源信息
        source_info = f"摄像头 {video_source}" if isinstance(video_source,
                                                             int) else f"视频: {os.path.basename(video_source)}"
        self.status_label.setText(f"正在从 {source_info} 加载视频流...")

        # 从滑块获取置信度阈值
        conf_threshold = self.confidence_slider.value() / 100

        # 创建并启动视频检测线程
        try:
            self.video_thread = VideoDetectionThread(self.model, video_source, conf_threshold)
            self.video_thread.frame_signal.connect(self.update_original_frame)
            self.video_thread.detection_signal.connect(self.update_detection_frame)
            self.video_thread.error_signal.connect(self.handle_video_error)
            self.video_thread.fps_signal.connect(self.update_fps)
            self.video_thread.start()

            # 更新UI状态
            self.is_video_mode = True
            self.stop_video_button.setEnabled(True)
            self.camera_button.setEnabled(False)
            self.video_button.setEnabled(False)
            self.load_button.setEnabled(False)
            self.detect_button.setEnabled(False)
            self.save_button.setEnabled(False)

            self.status_label.setText(f"正在从 {source_info} 进行实时检测")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"启动视频检测失败: {str(e)}")

    def stop_video_detection(self):
        """停止视频检测"""
        if self.video_thread is not None:
            self.video_thread.stop()
            self.video_thread = None

        # 恢复UI状态
        self.is_video_mode = False
        self.stop_video_button.setEnabled(False)
        self.camera_button.setEnabled(True)
        self.video_button.setEnabled(True)
        self.load_button.setEnabled(True)
        if self.current_image is not None:
            self.detect_button.setEnabled(True)

        self.fps_label.setText("FPS: 0.0")
        self.status_label.setText("视频检测已停止")

    def update_original_frame(self, frame):
        """更新原始视频帧"""
        if not self.is_video_mode:
            return

        self.current_image = frame

    def update_detection_frame(self, result):
        """更新检测结果帧和信息"""
        if not self.is_video_mode:
            return

        result_image, detections = result

        # 显示检测结果图像
        self.display_image(result_image)

        # 更新检测结果信息
        if not detections:
            self.results_label.setText("未检测到缺陷")
        else:
            result_text = f"<h3>检测到 {len(detections)} 个缺陷:</h3>"
            for i, det in enumerate(detections[:5]):  # 只显示前5个结果，避免刷新过快
                class_name = det['class_name']
                confidence = det['confidence']
                result_text += f"<p><b>缺陷 {i + 1}:</b> {class_name} ({confidence:.2f})</p>"

            if len(detections) > 5:
                result_text += f"<p>... 等 {len(detections) - 5} 个更多缺陷</p>"

            self.results_label.setText(result_text)

    def handle_video_error(self, error_message):
        """处理视频检测错误"""
        QMessageBox.warning(self, "视频检测警告", error_message)
        self.stop_video_detection()

    def update_fps(self, fps):
        """更新FPS显示"""
        self.fps_label.setText(f"FPS: {fps:.1f}")

    def save_results(self):
        """保存检测结果"""
        if self.current_image is None or self.detection_results is None:
            QMessageBox.warning(self, "警告", "没有可保存的检测结果")
            return

        # 保存图像
        options = QFileDialog.Options()
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存检测结果", "",
            "图像文件 (*.jpg *.png);;所有文件 (*)",
            options=options
        )

        if file_path:
            try:
                # 获取当前显示的图像
                pixmap = self.image_label.pixmap()
                if pixmap:
                    pixmap.save(file_path)
                    self.status_label.setText(f"结果已保存至: {file_path}")

                    # 保存检测结果文本
                    text_file = os.path.splitext(file_path)[0] + "_results.txt"
                    with open(text_file, 'w', encoding='utf-8') as f:
                        f.write(f"电路板缺陷检测结果报告\n")
                        f.write(f"图像文件: {os.path.basename(self.current_image_path)}\n")
                        f.write(f"检测时间: {import_time('datetime').datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")

                        if not self.detection_results:
                            f.write("未检测到缺陷\n")
                        else:
                            f.write(f"检测到 {len(self.detection_results)} 个缺陷:\n\n")
                            for i, det in enumerate(self.detection_results):
                                f.write(f"缺陷 {i + 1}:\n")
                                f.write(f"  类型: {det['class_name']}\n")
                                f.write(f"  置信度: {det['confidence']:.2f}\n")
                                x1, y1, x2, y2 = det['bbox']
                                f.write(f"  位置: ({x1}, {y1}) 到 ({x2}, {y2})\n\n")

                    QMessageBox.information(self, "保存成功",
                                            f"检测结果已保存:\n图像: {file_path}\n文本报告: {text_file}")
            except Exception as e:
                QMessageBox.critical(self, "保存错误", f"无法保存结果: {str(e)}")


def import_time(module_name):
    """导入模块并返回"""
    if module_name == 'datetime':
        import datetime
        return datetime
    return None


if __name__ == "__main__":
    # 创建应用
    app = QApplication(sys.argv)

    # 设置样式
    app.setStyle("Fusion")

    # 检查依赖项
    try:
        import torch
        import cv2

        # 检查ultralytics包
        try:
            import ultralytics
        except ImportError:
            print("正在安装ultralytics包...")
            import subprocess

            subprocess.check_call([sys.executable, "-m", "pip", "install", "ultralytics"])
            print("ultralytics安装完成")
            import ultralytics
    except ImportError as e:
        print(f"缺少依赖项: {str(e)}")
        QMessageBox.critical(None, "错误", f"缺少必要的依赖项: {str(e)}\n请安装所需的库后再运行程序。")
        sys.exit(1)

    # 创建目录
    os.makedirs("weights", exist_ok=True)

    # 创建并显示主窗口
    main_window = PCBDefectDetectionApp()
    main_window.show()

    # 运行应用
    sys.exit(app.exec_())